import os
from collections import defaultdict
from Bio import SeqIO
from ete3 import NCBITaxa


my_path = os.getcwd()
virus_fasta = os.path.join(my_path, 'e5.viruses.faa')
tax_dict = defaultdict()
ncbi = NCBITaxa()
for seq_record in SeqIO.parse(virus_fasta, 'fasta'):
    seq_id = str(seq_record.id)
    taxid = int(seq_id.split('.')[0])
    if taxid not in tax_dict:
        a = ncbi.get_taxid_translator([taxid])
        if a:
            tax_name = a[taxid]
            b = ncbi.get_lineage(taxid)
            names = ncbi.get_taxid_translator(b)
            lineage = []
            for i in b:
                lineage.append(str(i))
            taxonomy = [names[taxid] for taxid in b]
        else:
            tax_name = '#'
            lineage = ['#']
            taxonomy = ['#']
        tax_dict[taxid] = [lineage, taxonomy, tax_name]
tmp_egg_taxid_file = os.path.join(my_path, 'egg_taxid.txt')
result_lines = ''
with open(tmp_egg_taxid_file, 'r') as f:
    for each_line in f.readlines():
        result_lines += each_line.strip() + '\n'
result_file = os.path.join(my_path, 'combine_egg_taxid.txt')
with open(result_file, 'w') as g:
    for m, n in tax_dict.items():
        l_line = ','.join(n[0])
        t_line = ','.join(n[1])
        result_lines += '{0}\t{1}\tvirus\t{2}\t{3}\n'.format(str(m), n[2],
                                                             t_line, l_line)
    g.write(result_lines)
